# Please install OpenAI SDK first: `pip3 install openai`
from openai import OpenAI


class DeepSeekAPI:
    def __init__(self, api_key, base_url="https://api.deepseek.com/v1"):
        """
        初始化DeepSeek API客户端

        Args:
            api_key (str): DeepSeek API密钥
            base_url (str): API基础URL，默认为DeepSeek官方地址
        """
        self.client = OpenAI(
            api_key=api_key,
            base_url=base_url
        )

    def chat(self, model_name, query, system_prompt="我是一个AI智能助手，我可以帮助你解决各种问题。"):
        """
        与DeepSeek模型进行对话

        Args:
            model_name (str): 模型名称 ("deepseek-chat" 或 "deepseek-reasoner")
            query (str): 用户查询内容
            system_prompt (str): 系统提示词

        Returns:
            dict: 包含模型回复内容的字典
        """
        response = self.client.chat.completions.create(
            model=model_name,
            messages=[
                {"role": "system", "content": system_prompt},
                {"role": "user", "content": query},
            ],
            stream=False
        )

        result = {}
        if model_name == "deepseek-reasoner":
            result["content"] = response.choices[0].message.content
            result["reasoning_content"] = response.choices[0].message.reasoning_content
        else:
            result["content"] = response.choices[0].message.content

        return result


# 使用示例
if __name__ == '__main__':
    # 创建API实例
    dp_api = DeepSeekAPI(api_key="sk-942e904b9f1c45a1976e82a534eccbf7")

    # 发起对话
    res = dp_api.chat(
        model_name="deepseek-reasoner",
        query="你是什么大模型"
    )

    print(res)
